/*
 * Copyright 2018- The Pixie Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 */

#include "src/stirling/source_connectors/socket_tracer/protocols/mongodb/parse.h"

#include <string>
#include <utility>
#include <vector>

#include "src/common/testing/testing.h"

namespace px {
namespace stirling {
namespace protocols {
namespace mongodb {

// clang-format off

constexpr uint8_t mongoDBNeedMoreHeaderData[] = {
    // message length (4 bytes)
    0x00, 0x00, 0x00, 0x0c,
    // request id
    0x82, 0xb7, 0x31, 0x44,
    // response to
    0x00, 0x00, 0x00, 0x00,
    // op code (missing a byte)
    0xdd, 0x07, 0x00,
};

constexpr uint8_t mongoDBNeedMoreData[] = {
    // message length (18 bytes)
    0x12, 0x00, 0x00, 0x00,
    // request id
    0x82, 0xb7, 0x31, 0x44,
    // response to
    0x00, 0x00, 0x00, 0x00,
    // op code
    0xdd, 0x07, 0x00, 0x00,
    // flag bits (missing byte)
    0x00,
};

constexpr uint8_t mongoDBInvalidType[] = {
    // message length (18 bytes)
    0x12, 0x00, 0x00, 0x00,
    // request id
    0x82, 0xb7, 0x31, 0x44,
    // response to
    0x00, 0x00, 0x00, 0x00,
    // op code (2010, does not exist)
    0xda, 0x07, 0x00, 0x00,
    // flag bits
    0x00, 0x00,
};

constexpr uint8_t mongoDBUnsupportedType[] = {
    // message length (18 bytes)
    0x12, 0x00, 0x00, 0x00,
    // request id
    0x82, 0xb7, 0x31, 0x44,
    // response to
    0x00, 0x00, 0x00, 0x00,
    // op code (2004, not supported in newer versions)
    0xd4, 0x07, 0x00, 0x00,
    // flag bits
    0x00, 0x00,
};

constexpr uint8_t mongoDBInvalidFlagBits[] = {
    // message length (45 bytes)
    0x2d, 0x00, 0x00, 0x00,
    // request id (917)
    0x95, 0x03, 0x00, 0x00,
    // response to (444)
    0xbc, 0x01, 0x00, 0x00,
    // op code (2013)
    0xdd, 0x07, 0x00, 0x00,
    // flag bits (bits 2, 15 are set)
    0x04, 0x80, 0x00, 0x00,
    // section 1
    0x00, 0x18, 0x00, 0x00, 0x00, 0x10, 0x6e, 0x00, 0x01, 0x00, 0x00,
    0x00, 0x01, 0x6f, 0x6b, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xf0, 0x3f, 0x00,
};

constexpr uint8_t mongoDBValidFlagBitsSet[] = {
    // message length (49 bytes)
    0x31, 0x00, 0x00, 0x00,
    // request id (917)
    0x95, 0x03, 0x00, 0x00,
    // response to (444)
    0xbc, 0x01, 0x00, 0x00,
    // op code (2013)
    0xdd, 0x07, 0x00, 0x00,
    // flag bits (checksum, more to come, exhaust allowed set)
    0x03, 0x00, 0x01, 0x00,
    // section 1 (1 kind byte, 24 section body bytes)
    0x00, 0x18, 0x00, 0x00, 0x00, 0x10, 0x6e, 0x00, 0x01, 0x00, 0x00,
    0x00, 0x01, 0x6f, 0x6b, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xf0, 0x3f, 0x00,
    // checksum bytes
    0x00, 0x00, 0x00, 0x00,
};

constexpr uint8_t mongoDBMissingChecksum[] = {
    // message length (161 bytes)
    0x9d, 0x00, 0x00, 0x00,
    // request id (1144108930)
    0x82, 0xb7, 0x31, 0x44,
    // response to (0)
    0x00, 0x00, 0x00, 0x00,
    // op code (2013)
    0xdd, 0x07, 0x00, 0x00,
    // flag bits (checksum set)
    0x01, 0x00, 0x00, 0x00,
    // section 1 (1 kind byte, 82 section body bytes)
    0x00, 0x52, 0x00, 0x00, 0x00, 0x02, 0x69, 0x6e, 0x73, 0x65, 0x72,
    0x74, 0x00, 0x04, 0x00, 0x00, 0x00, 0x63, 0x61, 0x72, 0x00, 0x08,
    0x6f, 0x72, 0x64, 0x65, 0x72, 0x65, 0x64, 0x00, 0x01, 0x03, 0x6c,
    0x73, 0x69, 0x64, 0x00, 0x1e, 0x00, 0x00, 0x00, 0x05, 0x69, 0x64,
    0x00, 0x10, 0x00, 0x00, 0x00, 0x04, 0x0e, 0xab, 0xf5, 0xe5, 0x45,
    0xf8, 0x42, 0x5f, 0x8c, 0xb5, 0xb4, 0x0d, 0xff, 0x94, 0x8e, 0x1c,
    0x00, 0x02, 0x24, 0x64, 0x62, 0x00, 0x06, 0x00, 0x00, 0x00, 0x6d,
    0x79, 0x64, 0x62, 0x31, 0x00, 0x00,
    // section 2 (1 kind byte, 53 section body bytes)
    0x01, 0x35, 0x00, 0x00, 0x00, 0x64, 0x6f, 0x63, 0x75, 0x6d, 0x65,
    0x6e, 0x74, 0x73, 0x00, 0x27, 0x00, 0x00, 0x00, 0x07, 0x5f, 0x69,
    0x64, 0x00, 0x64, 0xdb, 0xd4, 0x67, 0x8f, 0x0e, 0x65, 0x5d, 0x43,
    0x14, 0xd6, 0x8a, 0x02, 0x6e, 0x61, 0x6d, 0x65, 0x00, 0x07, 0x00,
    0x00, 0x00, 0x74, 0x65, 0x73, 0x6c, 0x61, 0x34, 0x00, 0x00,
    // no checksum bytes
};

constexpr uint8_t mongoDBInvalidKindByte[] = {
    // message length (178 bytes)
    0xb2, 0x00, 0x00, 0x00,
    // request id (444)
    0xbc, 0x01, 0x00, 0x00,
    // response to (0)
    0x00, 0x00, 0x00, 0x00,
    // op code (2013)
    0xdd, 0x07, 0x00, 0x00,
    // flag bits
    0x00, 0x00, 0x00, 0x00,
    // section 1 (invalid kind byte)
    0x05, 0x9d, 0x00, 0x00, 0x00, 0x02, 0x69, 0x6e, 0x73, 0x65, 0x72,
    0x74, 0x00, 0x04, 0x00, 0x00, 0x00, 0x63, 0x61, 0x72, 0x00, 0x04,
    0x64, 0x6f, 0x63, 0x75, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x00, 0x40,
    0x00, 0x00, 0x00, 0x03, 0x30, 0x00, 0x38, 0x00, 0x00, 0x00, 0x02,
    0x6e, 0x61, 0x6d, 0x65, 0x00, 0x18, 0x00, 0x00, 0x00, 0x70, 0x69,
    0x78, 0x69, 0x65, 0x2d, 0x63, 0x61, 0x72, 0x2d, 0x31, 0x30, 0x2d,
    0x76, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x37, 0x2e, 0x30, 0x00,
    0x07, 0x5f, 0x69, 0x64, 0x00, 0x64, 0xe6, 0x72, 0x9c, 0x99, 0x6d,
    0x67, 0x6b, 0xf5, 0x20, 0x9d, 0xba, 0x00, 0x00, 0x08, 0x6f, 0x72,
    0x64, 0x65, 0x72, 0x65, 0x64, 0x00, 0x01, 0x03, 0x6c, 0x73, 0x69,
    0x64, 0x00, 0x1e, 0x00, 0x00, 0x00, 0x05, 0x69, 0x64, 0x00, 0x10,
    0x00, 0x00, 0x00, 0x04, 0xe7, 0xd7, 0x16, 0xb3, 0x75, 0xb7, 0x4c,
    0x39, 0x8b, 0x75, 0x41, 0x97, 0xc4, 0x97, 0x06, 0xd1, 0x00, 0x02,
    0x24, 0x64, 0x62, 0x00, 0x06, 0x00, 0x00, 0x00, 0x6d, 0x79, 0x64,
    0x62, 0x31, 0x00, 0x00,
};

constexpr uint8_t mongoDBInvalidSeqIdentifier[] = {
    // message length (157 bytes)
    0x9d, 0x00, 0x00, 0x00,
    // request id (1144108930)
    0x82, 0xb7, 0x31, 0x44,
    // response to (0)
    0x00, 0x00, 0x00, 0x00,
    // op code (2013)
    0xdd, 0x07, 0x00, 0x00,
    // flag bits
    0x00, 0x00, 0x00, 0x00,
    // section 1 (1 kind byte, 82 section body bytes)
    0x00, 0x52, 0x00, 0x00, 0x00, 0x02, 0x69, 0x6e, 0x73, 0x65, 0x72,
    0x74, 0x00, 0x04, 0x00, 0x00, 0x00, 0x63, 0x61, 0x72, 0x00, 0x08,
    0x6f, 0x72, 0x64, 0x65, 0x72, 0x65, 0x64, 0x00, 0x01, 0x03, 0x6c,
    0x73, 0x69, 0x64, 0x00, 0x1e, 0x00, 0x00, 0x00, 0x05, 0x69, 0x64,
    0x00, 0x10, 0x00, 0x00, 0x00, 0x04, 0x0e, 0xab, 0xf5, 0xe5, 0x45,
    0xf8, 0x42, 0x5f, 0x8c, 0xb5, 0xb4, 0x0d, 0xff, 0x94, 0x8e, 0x1c,
    0x00, 0x02, 0x24, 0x64, 0x62, 0x00, 0x06, 0x00, 0x00, 0x00, 0x6d,
    0x79, 0x64, 0x62, 0x31, 0x00, 0x00,
    // section 2 (invalid sequence identifier)
    0x01, 0x35, 0x00, 0x00, 0x00, 0x64, 0x6f, 0x63, 0x75, 0x6d, 0x65,
    0x6e, 0x74, 0x71, 0x00, 0x27, 0x00, 0x00, 0x00, 0x07, 0x5f, 0x69,
    0x64, 0x00, 0x64, 0xdb, 0xd4, 0x67, 0x8f, 0x0e, 0x65, 0x5d, 0x43,
    0x14, 0xd6, 0x8a, 0x02, 0x6e, 0x61, 0x6d, 0x65, 0x00, 0x07, 0x00,
    0x00, 0x00, 0x74, 0x65, 0x73, 0x6c, 0x61, 0x34, 0x00, 0x00,
};

constexpr uint8_t mongoDBEmptyDocument[] = {
    // message length (79 bytes)
    0x4f, 0x00, 0x00, 0x00,
    // request id (1144108930)
    0x82, 0xb7, 0x31, 0x44,
    // response to (0)
    0x00, 0x00, 0x00, 0x00,
    // op code (2013)
    0xdd, 0x07, 0x00, 0x00,
    // flag bits
    0x00, 0x00, 0x00, 0x00,
    // section 1 (1 kind byte, 0 section body bytes)
    0x00, 0x04, 0x00, 0x00, 0x00,
    // section 2 (1 kind byte, 53 section body bytes)
    0x01, 0x35, 0x00, 0x00, 0x00, 0x64, 0x6f, 0x63, 0x75, 0x6d, 0x65,
    0x6e, 0x74, 0x73, 0x00, 0x27, 0x00, 0x00, 0x00, 0x07, 0x5f, 0x69,
    0x64, 0x00, 0x64, 0xdb, 0xd4, 0x67, 0x8f, 0x0e, 0x65, 0x5d, 0x43,
    0x14, 0xd6, 0x8a, 0x02, 0x6e, 0x61, 0x6d, 0x65, 0x00, 0x07, 0x00,
    0x00, 0x00, 0x74, 0x65, 0x73, 0x6c, 0x61, 0x34, 0x00, 0x00,
};

constexpr uint8_t mongoDBValidRequest[] = {
    // message length (178 bytes)
    0xb2, 0x00, 0x00, 0x00,
    // request id (444)
    0xbc, 0x01, 0x00, 0x00,
    // response to (0)
    0x00, 0x00, 0x00, 0x00,
    // op code (2013)
    0xdd, 0x07, 0x00, 0x00,
    // flag bits
    0x00, 0x00, 0x00, 0x00,
    // section 1
    0x00, 0x9d, 0x00, 0x00, 0x00, 0x02, 0x69, 0x6e, 0x73, 0x65, 0x72,
    0x74, 0x00, 0x04, 0x00, 0x00, 0x00, 0x63, 0x61, 0x72, 0x00, 0x04,
    0x64, 0x6f, 0x63, 0x75, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x00, 0x40,
    0x00, 0x00, 0x00, 0x03, 0x30, 0x00, 0x38, 0x00, 0x00, 0x00, 0x02,
    0x6e, 0x61, 0x6d, 0x65, 0x00, 0x18, 0x00, 0x00, 0x00, 0x70, 0x69,
    0x78, 0x69, 0x65, 0x2d, 0x63, 0x61, 0x72, 0x2d, 0x31, 0x30, 0x2d,
    0x76, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x37, 0x2e, 0x30, 0x00,
    0x07, 0x5f, 0x69, 0x64, 0x00, 0x64, 0xe6, 0x72, 0x9c, 0x99, 0x6d,
    0x67, 0x6b, 0xf5, 0x20, 0x9d, 0xba, 0x00, 0x00, 0x08, 0x6f, 0x72,
    0x64, 0x65, 0x72, 0x65, 0x64, 0x00, 0x01, 0x03, 0x6c, 0x73, 0x69,
    0x64, 0x00, 0x1e, 0x00, 0x00, 0x00, 0x05, 0x69, 0x64, 0x00, 0x10,
    0x00, 0x00, 0x00, 0x04, 0xe7, 0xd7, 0x16, 0xb3, 0x75, 0xb7, 0x4c,
    0x39, 0x8b, 0x75, 0x41, 0x97, 0xc4, 0x97, 0x06, 0xd1, 0x00, 0x02,
    0x24, 0x64, 0x62, 0x00, 0x06, 0x00, 0x00, 0x00, 0x6d, 0x79, 0x64,
    0x62, 0x31, 0x00, 0x00,
};

constexpr uint8_t mongoDBValidResponse[] = {
    // message length (45 bytes)
    0x2d, 0x00, 0x00, 0x00,
    // request id (917)
    0x95, 0x03, 0x00, 0x00,
    // response to (444)
    0xbc, 0x01, 0x00, 0x00,
    // op code (2013)
    0xdd, 0x07, 0x00, 0x00,
    // flag bits
    0x00, 0x00, 0x00, 0x00,
    // section 1
    0x00, 0x18, 0x00, 0x00, 0x00, 0x10, 0x6e, 0x00, 0x01, 0x00, 0x00,
    0x00, 0x01, 0x6f, 0x6b, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xf0, 0x3f, 0x00,
};

constexpr uint8_t mongoDBValidRequestTwoSections[] = {
    // message length (157 bytes)
    0x9d, 0x00, 0x00, 0x00,
    // request id (1144108930)
    0x82, 0xb7, 0x31, 0x44,
    // response to (0)
    0x00, 0x00, 0x00, 0x00,
    // op code (2013)
    0xdd, 0x07, 0x00, 0x00,
    // flag bits
    0x00, 0x00, 0x00, 0x00,
    // section 1 (1 kind byte, 82 section body bytes)
    0x00, 0x52, 0x00, 0x00, 0x00, 0x02, 0x69, 0x6e, 0x73, 0x65, 0x72,
    0x74, 0x00, 0x04, 0x00, 0x00, 0x00, 0x63, 0x61, 0x72, 0x00, 0x08,
    0x6f, 0x72, 0x64, 0x65, 0x72, 0x65, 0x64, 0x00, 0x01, 0x03, 0x6c,
    0x73, 0x69, 0x64, 0x00, 0x1e, 0x00, 0x00, 0x00, 0x05, 0x69, 0x64,
    0x00, 0x10, 0x00, 0x00, 0x00, 0x04, 0x0e, 0xab, 0xf5, 0xe5, 0x45,
    0xf8, 0x42, 0x5f, 0x8c, 0xb5, 0xb4, 0x0d, 0xff, 0x94, 0x8e, 0x1c,
    0x00, 0x02, 0x24, 0x64, 0x62, 0x00, 0x06, 0x00, 0x00, 0x00, 0x6d,
    0x79, 0x64, 0x62, 0x31, 0x00, 0x00,
    // section 2 (1 kind byte, 53 section body bytes)
    0x01, 0x35, 0x00, 0x00, 0x00, 0x64, 0x6f, 0x63, 0x75, 0x6d, 0x65,
    0x6e, 0x74, 0x73, 0x00, 0x27, 0x00, 0x00, 0x00, 0x07, 0x5f, 0x69,
    0x64, 0x00, 0x64, 0xdb, 0xd4, 0x67, 0x8f, 0x0e, 0x65, 0x5d, 0x43,
    0x14, 0xd6, 0x8a, 0x02, 0x6e, 0x61, 0x6d, 0x65, 0x00, 0x07, 0x00,
    0x00, 0x00, 0x74, 0x65, 0x73, 0x6c, 0x61, 0x34, 0x00, 0x00,
};

constexpr uint8_t mongoDBValidResponseTwoSections[] = {
    // message length (45 bytes)
    0x2d, 0x00, 0x00, 0x00,
    // request id (444)
    0xbc, 0x01, 0x00, 0x00,
    // response to (1144108930)
    0x82, 0xb7, 0x31, 0x44,
    // op code (2013)
    0xdd, 0x07, 0x00, 0x00,
    // flag bits
    0x00, 0x00, 0x00, 0x00,
    // section 1 (1 kind byte, 24 section body bytes)
    0x00, 0x18, 0x00, 0x00, 0x00, 0x10, 0x6e, 0x00, 0x01, 0x00, 0x00,
    0x00, 0x01, 0x6f, 0x6b, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xf0, 0x3f, 0x00,
};

constexpr uint8_t mongoDBValidRequestAndInvalidRequest[] = {
    // valid frame
    // message length (45 bytes)
    0x2d, 0x00, 0x00, 0x00,
    // request id (917)
    0x95, 0x03, 0x00, 0x00,
    // response to (444)
    0xbc, 0x01, 0x00, 0x00,
    // op code (2013)
    0xdd, 0x07, 0x00, 0x00,
    // flag bits
    0x00, 0x00, 0x00, 0x00,
    // section 1
    0x00, 0x18, 0x00, 0x00, 0x00, 0x10, 0x6e, 0x00, 0x01, 0x00, 0x00,
    0x00, 0x01, 0x6f, 0x6b, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0xf0, 0x3f, 0x00,

    // invalid frame
    // message length (18 bytes)
    0x12, 0x00, 0x00, 0x00,
    // request id (444)
    0xbc, 0x01, 0x00, 0x00,
    // response to (0)
    0x00, 0x00, 0x00, 0x00,
    // flag bits
    0xFF, 0xFF, 0xFF, 0xFF,
    // section data
    0x00, 0x00
};

// clang-format on

class MongoDBParserTest : public ::testing::Test {};

TEST_F(MongoDBParserTest, ParseFrameWhenNeedsMoreHeaderData) {
  auto frame_view = CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBNeedMoreHeaderData));

  mongodb::Frame frame;
  StateWrapper state_order{};

  ParseState state = ParseFrame(message_type_t::kRequest, &frame_view, &frame, &state_order);

  EXPECT_EQ(state, ParseState::kNeedsMoreData);
}

TEST_F(MongoDBParserTest, ParseFrameWhenNeedsMoreData) {
  auto frame_view = CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBNeedMoreData));

  mongodb::Frame frame;
  StateWrapper state_order{};

  ParseState state = ParseFrame(message_type_t::kRequest, &frame_view, &frame, &state_order);

  EXPECT_EQ(state, ParseState::kNeedsMoreData);
}

TEST_F(MongoDBParserTest, ParseFrameWhenNotValidType) {
  auto frame_view = CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBInvalidType));

  mongodb::Frame frame;
  StateWrapper state_order{};

  ParseState state = ParseFrame(message_type_t::kRequest, &frame_view, &frame, &state_order);

  EXPECT_EQ(state, ParseState::kInvalid);
}

TEST_F(MongoDBParserTest, ParseFrameWhenUnsupportedType) {
  auto frame_view = CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBUnsupportedType));

  mongodb::Frame frame;
  StateWrapper state_order{};

  ParseState state = ParseFrame(message_type_t::kRequest, &frame_view, &frame, &state_order);

  EXPECT_EQ(state, ParseState::kIgnored);
  EXPECT_EQ(frame_view.length(), 0);
}

TEST_F(MongoDBParserTest, ParseFrameInvalidFlagBits) {
  auto frame_view = CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBInvalidFlagBits));

  mongodb::Frame frame;
  StateWrapper state_order{};

  ParseState state = ParseFrame(message_type_t::kRequest, &frame_view, &frame, &state_order);

  EXPECT_EQ(state, ParseState::kInvalid);
}

TEST_F(MongoDBParserTest, ParseFrameValidFlagBitsSet) {
  auto frame_view = CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBValidFlagBitsSet));

  mongodb::Frame frame;
  StateWrapper state_order{};

  ParseState state = ParseFrame(message_type_t::kRequest, &frame_view, &frame, &state_order);

  EXPECT_TRUE(frame.checksum_present);
  EXPECT_TRUE(frame.more_to_come);
  EXPECT_TRUE(frame.exhaust_allowed);
  EXPECT_EQ(state, ParseState::kSuccess);
}

TEST_F(MongoDBParserTest, ParseFrameInvalidChecksum) {
  auto frame_view = CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBMissingChecksum));

  mongodb::Frame frame;
  StateWrapper state_order{};

  ParseState state = ParseFrame(message_type_t::kRequest, &frame_view, &frame, &state_order);

  EXPECT_TRUE(frame.checksum_present);
  EXPECT_EQ(state, ParseState::kInvalid);
}

TEST_F(MongoDBParserTest, ParseFrameInvalidKindByte) {
  auto frame_view = CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBInvalidKindByte));

  mongodb::Frame frame;
  StateWrapper state_order{};

  ParseState state = ParseFrame(message_type_t::kRequest, &frame_view, &frame, &state_order);

  EXPECT_EQ(state, ParseState::kInvalid);
}

TEST_F(MongoDBParserTest, ParseFrameInvalidSeqIdentifier) {
  auto frame_view =
      CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBInvalidSeqIdentifier));

  mongodb::Frame frame;
  StateWrapper state_order{};

  ParseState state = ParseFrame(message_type_t::kRequest, &frame_view, &frame, &state_order);

  EXPECT_EQ(state, ParseState::kInvalid);
}

TEST_F(MongoDBParserTest, ParseFrameEmptyDocument) {
  auto frame_view = CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBEmptyDocument));

  mongodb::Frame frame;
  StateWrapper state_order{};

  ParseState state = ParseFrame(message_type_t::kRequest, &frame_view, &frame, &state_order);

  EXPECT_EQ(frame.sections[0].kind, 0);
  EXPECT_EQ(frame.sections[0].length, 4);
  EXPECT_EQ(frame.sections[0].documents[0].length(), 0);
  EXPECT_EQ(frame.sections[1].kind, 1);
  EXPECT_EQ(frame.sections[1].length, 53);
  EXPECT_EQ(frame.sections[1].documents[0].length(), 70);
  EXPECT_EQ(state, ParseState::kSuccess);
}

TEST_F(MongoDBParserTest, ParseFrameValidRequest) {
  auto frame_view = CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBValidRequest));

  mongodb::Frame frame;
  StateWrapper state_order{};

  ParseState state = ParseFrame(message_type_t::kRequest, &frame_view, &frame, &state_order);

  EXPECT_EQ(frame.length, 178);
  EXPECT_EQ(frame.request_id, 444);
  EXPECT_EQ(frame.response_to, 0);
  EXPECT_EQ(frame.op_code, static_cast<int32_t>(mongodb::Type::kOPMsg));
  EXPECT_EQ(frame.sections[0].kind, 0);
  EXPECT_EQ(frame.sections[0].length, 157);
  EXPECT_EQ(frame.op_msg_type, "insert");
  EXPECT_EQ(state, ParseState::kSuccess);
}

TEST_F(MongoDBParserTest, ParseFrameValidResponse) {
  auto frame_view = CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBValidResponse));

  mongodb::Frame frame;
  StateWrapper state_order{};

  ParseState state = ParseFrame(message_type_t::kResponse, &frame_view, &frame, &state_order);

  EXPECT_EQ(frame.length, 45);
  EXPECT_EQ(frame.request_id, 917);
  EXPECT_EQ(frame.response_to, 444);
  EXPECT_EQ(frame.op_code, static_cast<int32_t>(mongodb::Type::kOPMsg));
  EXPECT_EQ(frame.sections[0].kind, 0);
  EXPECT_EQ(frame.sections[0].length, 24);
  EXPECT_EQ(frame.op_msg_type, "ok: {$numberDouble: 1.0}");
  EXPECT_EQ(state, ParseState::kSuccess);
}

TEST_F(MongoDBParserTest, ParseFrameValidRequestTwoSections) {
  auto frame_view =
      CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBValidRequestTwoSections));

  mongodb::Frame frame;
  StateWrapper state_order{};

  ParseState state = ParseFrame(message_type_t::kRequest, &frame_view, &frame, &state_order);

  EXPECT_EQ(frame.length, 157);
  EXPECT_EQ(frame.request_id, 1144108930);
  EXPECT_EQ(frame.response_to, 0);
  EXPECT_EQ(frame.op_code, static_cast<int32_t>(mongodb::Type::kOPMsg));
  EXPECT_EQ(frame.sections[0].kind, 0);
  EXPECT_EQ(frame.sections[0].length, 82);
  EXPECT_EQ(frame.op_msg_type, "insert");
  EXPECT_EQ(frame.sections[1].kind, 1);
  EXPECT_EQ(frame.sections[1].length, 53);
  EXPECT_EQ(state, ParseState::kSuccess);
}

TEST_F(MongoDBParserTest, ParseFrameValidResponseTwoSections) {
  auto frame_view =
      CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBValidResponseTwoSections));

  mongodb::Frame frame;
  StateWrapper state_order{};

  ParseState state = ParseFrame(message_type_t::kResponse, &frame_view, &frame, &state_order);

  EXPECT_EQ(frame.length, 45);
  EXPECT_EQ(frame.request_id, 444);
  EXPECT_EQ(frame.response_to, 1144108930);
  EXPECT_EQ(frame.op_code, static_cast<int32_t>(mongodb::Type::kOPMsg));
  EXPECT_EQ(frame.sections[0].kind, 0);
  EXPECT_EQ(frame.sections[0].length, 24);
  EXPECT_EQ(frame.op_msg_type, "ok: {$numberDouble: 1.0}");
  EXPECT_EQ(state, ParseState::kSuccess);
}

TEST_F(MongoDBParserTest, ParseValidFrameAndInvalidFrame) {
  auto frame_view =
      CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBValidRequestAndInvalidRequest));

  mongodb::Frame valid_frame;
  StateWrapper state_order{};

  ParseState state = ParseFrame(message_type_t::kRequest, &frame_view, &valid_frame, &state_order);

  EXPECT_EQ(valid_frame.length, 45);
  EXPECT_EQ(valid_frame.request_id, 917);
  EXPECT_EQ(valid_frame.response_to, 444);
  EXPECT_EQ(valid_frame.op_code, static_cast<int32_t>(mongodb::Type::kOPMsg));
  EXPECT_EQ(valid_frame.sections[0].kind, 0);
  EXPECT_EQ(valid_frame.sections[0].length, 24);
  EXPECT_EQ(valid_frame.op_msg_type, "ok: {$numberDouble: 1.0}");
  EXPECT_EQ(state, ParseState::kSuccess);

  mongodb::Frame invalid_frame;
  state = ParseFrame(message_type_t::kRequest, &frame_view, &invalid_frame, &state_order);
  EXPECT_EQ(state, ParseState::kInvalid);
}

TEST_F(MongoDBParserTest, ValidateStateOrderFromFrames) {
  // Setup three different request frames.
  auto frame_view_1 = CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBValidFlagBitsSet));
  auto frame_view_2 = CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBValidRequest));
  auto frame_view_3 =
      CreateStringView<char>(CharArrayStringView<uint8_t>(mongoDBValidRequestTwoSections));

  // Initialize what the state's stream order should look like after the first insertion.
  std::vector<std::pair<mongodb::stream_id_t, bool>> stream_order{{917, false}};

  StateWrapper state_order{};

  mongodb::Frame frame_1;
  ParseState state_1 = ParseFrame(message_type_t::kRequest, &frame_view_1, &frame_1, &state_order);
  EXPECT_EQ(frame_1.request_id, 917);
  EXPECT_EQ(state_order.global.stream_order, stream_order);
  EXPECT_EQ(state_1, ParseState::kSuccess);

  stream_order.push_back({444, false});
  mongodb::Frame frame_2;
  ParseState state_2 = ParseFrame(message_type_t::kRequest, &frame_view_2, &frame_2, &state_order);
  EXPECT_EQ(frame_2.request_id, 444);
  EXPECT_EQ(state_order.global.stream_order, stream_order);
  EXPECT_EQ(state_2, ParseState::kSuccess);

  stream_order.push_back({1144108930, false});
  mongodb::Frame frame_3;
  ParseState state_3 = ParseFrame(message_type_t::kRequest, &frame_view_3, &frame_3, &state_order);
  EXPECT_EQ(frame_3.request_id, 1144108930);
  EXPECT_EQ(state_order.global.stream_order, stream_order);
  EXPECT_EQ(state_3, ParseState::kSuccess);
}

}  // namespace mongodb
}  // namespace protocols
}  // namespace stirling
}  // namespace px
